import os
import sys
import glob
import time
import pickle
from pathlib import Path

import qlib
from qlib.constant import REG_CN
from qlib.workflow import R

BASE = Path(__file__).resolve().parents[1] / 'mlruns'


def read_last_metric_value(metric_path: Path):
    try:
        # mlflow metrics files format: timestamp value step\n
        lines = metric_path.read_text(encoding='utf-8', errors='ignore').strip().splitlines()
        if not lines:
            return None
        last = lines[-1].strip().split()
        if len(last) >= 2:
            return float(last[1])
    except Exception:
        return None
    return None


def find_best_run(base_dir: Path):
    best = None
    # search across all experiments under mlruns, skipping .trash
    for exp_dir in base_dir.iterdir():
        if not exp_dir.is_dir() or exp_dir.name == '.trash':
            continue
        for run_dir in exp_dir.iterdir():
            if not run_dir.is_dir():
                continue
            metrics_dir = run_dir / 'metrics'
            artifacts_dir = run_dir / 'artifacts'
            if not metrics_dir.exists() and not artifacts_dir.exists():
                continue
            ic = read_last_metric_value(metrics_dir / 'IC') if metrics_dir.exists() else None
            ric = read_last_metric_value(metrics_dir / 'Rank IC') if metrics_dir.exists() else None
            score = None
            if ric is not None:
                score = ric
            elif ic is not None:
                score = ic
            else:
                if (artifacts_dir / 'pred.pkl').exists():
                    score = run_dir.stat().st_mtime
            if score is None:
                continue
            if best is None or score > best['score']:
                best = {
                    'run_id': run_dir.name,
                    'score': score,
                    'ic': ic,
                    'ric': ric,
                    'artifacts': artifacts_dir,
                    'metrics': metrics_dir,
                    'experiment': exp_dir.name,
                }
    return best


def main():
    provider = os.path.expanduser('~/.qlib/qlib_data/cn_data')
    qlib.init(provider_uri=provider, region=REG_CN)

    if not BASE.exists():
        print('mlruns path not found:', BASE)
        sys.exit(1)

    best = find_best_run(BASE)
    if not best:
        print('No suitable run found under', BASE)
        sys.exit(2)

    run_id = best['run_id']
    print(f"Selected RUN_ID: {run_id} (exp {best['experiment']}) | Rank IC={best['ric']} | IC={best['ic']}")

    # Try to load model and dataset for fresh prediction
    model = None
    dataset = None
    try:
        # Directly load artifacts from disk to avoid recorder lookup issues
        art_dir = BASE / best['experiment'] / run_id / 'artifacts'
        params_pkl = art_dir / 'params.pkl'
        dataset_pkl = art_dir / 'dataset'
        # prefer 'params.pkl' then 'model.pkl'
        if params_pkl.exists():
            with open(params_pkl, 'rb') as f:
                model = pickle.load(f)
        else:
            model = None
        if dataset_pkl.exists():
            with open(dataset_pkl, 'rb') as f:
                dataset = pickle.load(f)
        # Override test segment for demonstration if needed
        segments = getattr(dataset, 'segments', None)
        if segments and 'test' in segments:
            pass
        else:
            try:
                dataset.config(segments={'test': ('2019-01-01', '2020-08-01')})
            except Exception:
                pass
        # Predict
        pred_series = model.predict(dataset)
        try:
            import pandas as pd
            pred_df = pred_series.to_frame('score') if hasattr(pred_series, 'to_frame') else pd.DataFrame({'score': pred_series})
        except Exception:
            pred_df = pred_series
        print('Fresh prediction (head):')
        print(pred_df.head(10))
        out_path = BASE.parent / f'pred_best_{run_id}.pkl'
        with open(out_path, 'wb') as f:
            pickle.dump(pred_df, f)
        print('Saved new prediction to', out_path)
        return
    except Exception as e:
        print('Fresh prediction failed, falling back to existing pred.pkl. Reason:', repr(e))

    # Fallback: load existing prediction artifact from disk
    try:
        art_dir = BASE / best['experiment'] / run_id / 'artifacts'
        pred_pkl = art_dir / 'pred.pkl'
        with open(pred_pkl, 'rb') as f:
            pred = pickle.load(f)
        print('Existing prediction (head):')
        try:
            print(pred.head(10))
        except Exception:
            print(pred)
        out_path = BASE.parent / f'pred_best_{run_id}.pkl'
        with open(out_path, 'wb') as f:
            pickle.dump(pred, f)
        print('Saved existing prediction copy to', out_path)
    except Exception as e:
        print('Failed to load existing pred.pkl:', repr(e))
        sys.exit(3)


if __name__ == '__main__':
    main()